#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Keyword, DML
import pymysql
from sqlalchemy import create_engine
import urllib
from DBUtils.PooledDB import PooledDB

# https://groups.google.com/forum/#!forum/sqlparse/browse_thread/thread/b0bd9a022e9d4895
################################################################################
def is_subselect(parsed):
    if not parsed.is_group:
        return False
    for item in parsed.tokens:
        if item.ttype is DML and item.value.upper() == 'SELECT':
            return True
    return False


def extract_from_part(parsed):
    from_seen = False
    for item in parsed.tokens:
        if from_seen:
            if is_subselect(item):
                for x in extract_from_part(item):
                    yield x
            elif item.ttype is Keyword:
                return
            else:
                yield item
        elif item.ttype is Keyword and item.value.upper() == 'FROM':
            from_seen = True


def extract_table_identifiers(token_stream):
    ''' 提取表名和别名
    print('table name:%s, table alias:%s, table real name:%s' %
    (tables[0].get_name(), tables[0].get_alias(), tables[0].get_real_name()))
    '''
    for item in token_stream:
        if isinstance(item, IdentifierList):
            for identifier in item.get_identifiers():
                yield identifier.get_real_name()
        elif isinstance(item, Identifier):
            yield item.get_real_name()
        # It's a bug to check for Keyword here, but in the example
        # above some tables names are identified as keywords...
        elif item.ttype is Keyword:
            yield item.value


def extract_tables(sql):
    """
    Extract the table names from an SQL statment.
    Returns a list of TableReference namedtuples.
    """
    stream = extract_from_part(sqlparse.parse(sql)[0])
    return list(extract_table_identifiers(stream))


def extract_sql_columns(sql):
    columns = []
    parsed = sqlparse.parse(sql)
    stmt = parsed[0]
    for token in stmt.tokens:
        if isinstance(token, IdentifierList):
            for identifier in token.get_identifiers():
                columns.append(identifier.get_real_name())
        if isinstance(token, Identifier):
            columns.append(token.get_real_name())
        if token.ttype is Keyword:  # from
            break
    return columns
################################################################################

def optimizeFieldAccess(grammer, cursor):
    '''
    把基于名字的数据访问df['name']转变是基于顺序的数据访问df[n]，前者需要字典dict数据类型，后者只需要元组tuple或数组
    '''
    if isinstance(cursor, list):
        for i, field in enumerate(cursor):
            if isinstance(field, str):
                grammer = grammer.replace('df["'+field+'"]', 'df['+str(i) +']')
    else:
        for i, field in enumerate(cursor.description):
            grammer = grammer.replace('df["'+field[0]+'"]', 'df['+str(i) +']')
    return grammer


def getConn(host='127.0.0.1', user='root', password='', dbname='', dictResult=True):
    '''
    建立数据库连接
    '''
    try:
        if dictResult:
            cursorclass = pymysql.cursors.SSDictCursor
        else:
            cursorclass = pymysql.cursors.SSCursor
        conn = pymysql.connect(host = host,
                               user = user,
                               password = password,
                               db = dbname,
                               charset = 'utf8mb4',
                               cursorclass = cursorclass)
        return conn
    except:
        return None


def getDbpool(host='127.0.0.1', user='root', password='', dbname='', dictResult=True, url=''):
    '''
    建立连接池，返回连接池地址
    '''

    def getDbClassInfo(dbn):
        try:
            dbn_lower = dbn.lower()

            drivers = {
                'mysql': 'pymysql',
                'postgres': ['psycopg2', 'psycopg', 'pgdb'],
                'sqlite': ['sqlite3', 'pysqlite2.dbapi2', 'sqlite'],
                'hive': (['pyhive'], 'presto'),
            }[dbn_lower]

            if isinstance(drivers, tuple):
                fromlist = drivers[1]
                drivers = drivers[0]
            else:
                fromlist = []

            if (dbn_lower == 'pymysql'):
                # Mysql driver.
                #
                return pymysql, pymysql.cursors.SSDictCursor, pymysql.cursors.SSCursor

            elif (dbn_lower == 'postgres'):
                # Postgres driver.
                #
                db_module = import_driver(drivers, fromlist)
                if db_module.__name__ == "psycopg2":
                    import psycopg2.extensions

                    psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
                if db_module.__name__ == "pgdb" and 'port' in keywords:
                    keywords["host"] += ":" + str(keywords.pop('port'))
                return db_module, None, None

            elif (dbn_lower == 'sqlite'):
                # Sqlite driver.
                #
                db_module = import_driver(drivers, fromlist)
                if db_module.__name__ in ["sqlite3", "pysqlite2.dbapi2"]:
                    db_module.paramstyle = 'qmark'

                return db_module, None, None
        except KeyError:
            return (None, None, None)

    if url:
        dbinfo = dburl2dict(url)
        creator, cursordict, cursortuple = getDbClassInfo(dbn)

        kargs = {}
        if dictResult:
            kargs['cursorclass'] = cursordict
        else:
            kargs['cursorclass'] = cursortuple

        try:
            conn_pool = PooledDB( creator = creator,
                                  host = host,
                                  user = user,
                                  password = password,
                                  database = dbname,
                                  autocommit = True,
                                  charset = 'utf8mb4',
                                  **kargs)
                                  #cursorclass = cursorclass)
            return conn_pool
        except (BaseException) as exc:
            return str(exc)
    else:
        try:
            if dictResult:
                cursorclass = pymysql.cursors.SSDictCursor
            else:
                cursorclass = pymysql.cursors.SSCursor
            conn_pool = PooledDB( creator = pymysql,
                                  host = host,
                                  user = user,
                                  password = password,
                                  database = dbname,
                                  autocommit = True,
                                  charset = 'utf8mb4',
                                  cursorclass = cursorclass)
            return conn_pool
        except (BaseException) as exc:
            return str(exc)


def dburl2dict(url):
    """
    Takes a URL to a database and parses it into an equivalent dictionary.

        >>> dburl2dict('postgres://james:day@serverfarm.example.net:5432/mygreatdb')
        {'pw': 'day', 'dbn': 'postgres', 'db': 'mygreatdb', 'host': 'serverfarm.example.net', 'user': 'james', 'port': '5432'}
        >>> dburl2dict('postgres://james:day@serverfarm.example.net/mygreatdb')
        {'user': 'james', 'host': 'serverfarm.example.net', 'db': 'mygreatdb', 'pw': 'day', 'dbn': 'postgres'}
        >>> dburl2dict('postgres://james:d%40y@serverfarm.example.net/mygreatdb')
        {'user': 'james', 'host': 'serverfarm.example.net', 'db': 'mygreatdb', 'pw': 'd@y', 'dbn': 'postgres'}
    """
    dbn, rest = url.split('://', 1)
    user, rest = rest.split(':', 1)
    pw, rest = rest.split('@', 1)
    if ':' in rest:
        host, rest = rest.split(':', 1)
        port, rest = rest.split('/', 1)
    else:
        host, rest = rest.split('/', 1)
        port = None
    db = rest

    uq = urllib.unquote
    out = dict(dbms_name=dbn, db_user=uq(user), db_password=uq(pw), db_server=uq(host), database=uq(db))
    if port: out['port'] = port
    return out


def import_driver(drivers, fromlist=[]):
    """Import the first available driver.
    """
    for d in drivers:
        try:
            return __import__(d, None, None, fromlist)
        except ImportError:
            pass
    raise ImportError("Unable to import " + " or ".join(drivers))


def dictToSQLInsert(dict_, table_name):
    '''
    Using a Python dict for a SQL INSERT statement
    INSERT INTO table (dict.keys()) VALUES dict.values()
    https://stackoverflow.com/questions/9336270/using-a-python-dict-for-a-sql-insert-statement
    '''
    qmarks = ', '.join('?' * len(dict_))
    qry = "INSERT INTO ? (%s) VALUES (%s)" % (qmarks, qmarks)
    #cursor.execute(qry, list(dict_.keys()) + list(dict_.values()))
    return qry, [table_name] + list(dict_.keys()) + list(dict_.values())


def prepareSqlite3(db_url='sqlite:///:memory:'):
    '''
    https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#io-sql

    cursor.executescript("""drop table data""")
    df.to_sql('data', engine)
    pd.read_sql_table(table_name='data', con=engine)
    '''
    engine = create_engine(db_url)
    conn = engine.connect().connection
    cursor = conn.cursor()
    return engine, conn, cursor


if __name__ == '__main__':
    sql = """
    select K.a,K.b from (select H.b from (select G.c from (select F.d from
    (select E.e from A, B, C, D, E), F), G), H), I, J, K order by 1,2;
    """

    tables = ', '.join(extract_tables(sql))
    print('Tables: {0}'.format(tables))
